#include "DeferredProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "LocalShader.h"
#include "LocalScene.h"
#include "Math/Random.h"

void DeferredProcess::OnInit()
{
	mLocalScene = static_cast<LocalScene*>(mGameScene);

	mDebugIndex = 0;

	for (int i = 0; i < VPL_NUM; ++i)
	{
		float xi1 = Random64::Rand(0.f, 1.f);
		float xi2 = Random64::Rand(0.f, 1.f);

		mSampleCoordsAndWeights[i] = glm::vec4(xi1 * sin(2 * Math::M_PI * xi2), xi1 * cos(2 * Math::M_PI * xi2), xi1 * xi1, 0);
	}

	Ubo ubo;
	ubo.cameraView = mGameScene->GetCamera().GetView();
	ubo.lightProj = mLocalScene->mLightProj;
	ubo.lightView = mLocalScene->mLightView;
	ubo.lightData = glm::vec4(mLocalScene->mLightColor, mLocalScene->mLightIntensity);
	ubo.lightDir = mLocalScene->mLightDirection;
	ubo.debugIndex = mDebugIndex;

	ENQUEUE_RENDER_COMMAND(BuildModel)([this, ubo](VkCommandBuffer commandBuffer)
		{
			Ubo uboData = ubo;

			mUboBuffer = CustomBuffer::Builder()
				.SetData(&uboData)
				.SetSize(sizeof(Ubo))
				.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
				.Build();

			mSampleCoordsAndWeightsBuffer = CustomBuffer::Builder()
				.SetData(mSampleCoordsAndWeights.data())
				.SetSize(sizeof(glm::vec4) * mSampleCoordsAndWeights.size())
				.SetUsage(VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT)
				.Build();

			RefCountPtr<RHITexture> mrtPositionImage = GetTexture("MrtProcess", "positionRenderTarget");
			RefCountPtr<RHITexture> mrtNormalImage = GetTexture("MrtProcess", "normalRenderTarget");
			RefCountPtr<RHITexture> mrtAlbedoImage = GetTexture("MrtProcess", "albedoRenderTarget");

			RefCountPtr<RHITexture> rsmPositionImage = GetTexture("RsmProcess", "positionRenderTarget");
			RefCountPtr<RHITexture> rsmNormalImage = GetTexture("RsmProcess", "normalRenderTarget");
			RefCountPtr<RHITexture> rsmFluxImage = GetTexture("RsmProcess", "fluxRenderTarget");

			// mCommonSampler = RHISampler::Builder().Build();
			mCommonSampler = RHISampler::Builder()
				.SetMipmapMode(VK_SAMPLER_MIPMAP_MODE_LINEAR)
				.Build();

			mRsmSampler = RHISampler::Builder()
				.SetMipmapMode(VK_SAMPLER_MIPMAP_MODE_LINEAR)
				.SetBorderColor(VK_BORDER_COLOR_FLOAT_OPAQUE_BLACK)
				.SetAddressModeU(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_BORDER)
				.SetAddressModeV(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_BORDER)
				.SetAddressModeW(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_BORDER)
				.Build();

			mDescriptor = RHIDescriptor::Builder()
				.PushShader<DeferredVertexShader>()
				.PushShader<DeferredFragmentShader>()
				.SetDescriptorCount(0, 1)
				.Build();

			mDescriptor->SetValue("mrtPositionImage", mrtPositionImage, mCommonSampler);
			mDescriptor->SetValue("mrtNormalImage", mrtNormalImage, mCommonSampler);
			mDescriptor->SetValue("mrtAlbedoImage", mrtAlbedoImage, mCommonSampler);
			mDescriptor->SetValue("rsmPositionImage", rsmPositionImage, mRsmSampler);
			mDescriptor->SetValue("rsmNormalImage", rsmNormalImage, mRsmSampler);
			mDescriptor->SetValue("rsmFluxImage", rsmFluxImage, mRsmSampler);
			mDescriptor->SetValue("ubo", mUboBuffer);
			mDescriptor->SetValue("sampleCoordsAndWeights", mSampleCoordsAndWeightsBuffer);
			mDescriptorMask = mDescriptor->AllocateDescriptor();

			mColorRenderTarget = ColorRenderTarget::Builder()
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "colorRenderTarget");

			mDepthRenderTarget = DepthRenderTarget::Builder()
				.Build(this, "depthRenderTarget");

			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.AddColorRenderTarget(mColorRenderTarget)
				.SetDepthRenderTarget(mDepthRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.Build();

			mRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(mColorRenderTarget)
				.AddRenderTarget(mDepthRenderTarget)
				.AddPipeline(mPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build(this, "RenderPass");

			mFrameBuffer = RHIFrameBuffer::Builder()
				.SetRenderPass(mRenderPass)
				.Build();
		});
}

void DeferredProcess::OnTick(float deltaTime)
{
	Ubo ubo;
	ubo.cameraView = mGameScene->GetCamera().GetView();
	ubo.lightProj = mLocalScene->mLightProj;
	ubo.lightView = mLocalScene->mLightView;
	ubo.lightData = glm::vec4(mLocalScene->mLightColor, mLocalScene->mLightIntensity);
	ubo.lightDir = mLocalScene->mLightDirection;
	ubo.debugIndex = mDebugIndex;

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this, ubo](VkCommandBuffer commandBuffer)
		{
			Ubo uboData = ubo;

			mUboBuffer->WriteMemory(&uboData);

			VkRenderPassBeginInfo renderPassBeginInfo = mFrameBuffer->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(commandBuffer, 0, 1, &GRenderState->GetViewport());
			vkCmdSetScissor(commandBuffer, 0, 1, &GRenderState->GetScissor());

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 1, &mDescriptorMask.GetDescriptorSet(), 0, nullptr);

			vkCmdDraw(commandBuffer, 3, 1, 0, 0);

			vkCmdEndRenderPass(commandBuffer);
		});
}

void DeferredProcess::OnResize()
{
	RefCountPtr<RHITexture> mrtPositionImage = GetTexture("MrtProcess", "positionRenderTarget");
	RefCountPtr<RHITexture> mrtNormalImage = GetTexture("MrtProcess", "normalRenderTarget");
	RefCountPtr<RHITexture> mrtAlbedoImage = GetTexture("MrtProcess", "albedoRenderTarget");

	mDescriptor->SetValue("mrtPositionImage", mrtPositionImage, mCommonSampler);
	mDescriptor->SetValue("mrtNormalImage", mrtNormalImage, mCommonSampler);
	mDescriptor->SetValue("mrtAlbedoImage", mrtAlbedoImage, mCommonSampler);

	mDescriptor->UpdateDescriptor(mDescriptorMask);


	mColorRenderTarget = ColorRenderTarget::Builder()
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "colorRenderTarget");

	mDepthRenderTarget = DepthRenderTarget::Builder()
		.Build(this, "depthRenderTarget");

	mFrameBuffer = RHIFrameBuffer::Builder()
		.AddRenderTarget(mColorRenderTarget)
		.AddRenderTarget(mDepthRenderTarget)
		.SetRenderPass(mRenderPass)
		.Build();
}
